import pandas as pd


def dm_info(sql: str) -> pd.DataFrame:
    with open(sql, 'r', encoding='utf-8') as file:
        sql_content = file.read().split("DROP FUNCTION")[0].split("CREATE OR REPLACE FUNCTION")[0].split("ADD CONSTRAINT")[0]
    sql_split = [item.strip() for item in sql_content.split("CREATE TABLE \"KNOWLEDGE\".\"") if item and item.strip()]

    dm_table_field = pd.DataFrame(columns=['uid', '表名PG', '字段PG', '完整类型PG', '类型PG', '长度PG'])

    head = True
    for sql_one in sql_split:
        if head:
            head = False
            continue

        # print("=========================")

        tableName = sql_one.split("\"")[0]

        fieldArr = [item.strip() for item in sql_one.split("\n") if item and item.strip() and item.strip().startswith("\"")]
        for i in range(0, len(fieldArr)):
            fieldOneSplit = fieldArr[i].split(" ")
            fieldName = fieldOneSplit[0].replace("\"", "").strip()
            fieldTypeFull = fieldOneSplit[1].replace(",", "").strip()
            strip_split = fieldTypeFull.split("(")
            dm_table_field.loc[len(dm_table_field)] = [f"{tableName}.{fieldName}",
                                                       tableName,
                                                       fieldName,
                                                       fieldTypeFull,
                                                       strip_split[0],
                                                       strip_split[1].replace(")", "") if len(strip_split) > 1 else ""
                                                       ]

    dm_field_comment = pd.DataFrame(columns=['uid', '字段描述PG'])
    comment_arr = [item.strip() for item in sql_content.split("\n") if item and item.strip() and item.strip().startswith("COMMENT ON COLUMN \"KNOWLEDGE\".")]
    for comment_one in comment_arr:
        c_table = comment_one.split("\"")[3]
        c_column = comment_one.split("\"")[5]
        c_comment = comment_one.split("'")[1]
        # print(f"{c_table}\t{c_column}\t{c_comment}")
        dm_field_comment.loc[len(dm_field_comment)] = [f"{c_table}.{c_column}", c_comment]

    dm_table_comment = pd.DataFrame(columns=['表名PG', '表描述PG'])
    table_comment_arr = [item.strip() for item in sql_content.split("\n") if item and item.strip() and item.strip().startswith("COMMENT ON TABLE \"KNOWLEDGE\".")]
    for comment_one in table_comment_arr:
        print(comment_one)
        c_table = comment_one.split("\"")[3]
        c_comment = comment_one.split("'")[1]
        dm_table_comment.loc[len(dm_table_comment)] = [c_table, c_comment]

    dm_table_field = dm_table_field.merge(dm_field_comment, how='left', on=['uid'])
    dm_table_field = dm_table_field.merge(dm_table_comment, how='left', on=['表名PG'])
    return dm_table_field


if __name__ == '__main__':
    info = dm_info('pg.sql')
    info.to_excel("gen/高斯数据库.xlsx", index=False)
    print(info.head(100))
